#include <algorithm>
#include <chrono>
#include <cmath>
#include <fstream>
#include <iostream>
#include <limits>
#include <omp.h>
#include <random>
#include <vector>

// 将数据保存到txt文件
std::ofstream true_centers_file("true_centers.txt");
std::ofstream initial_data_file("initial_data.txt");
std::ofstream after_kmeans_centers_file("after_kmeans_centers.txt");
std::ofstream after_kmeans_omp_centers_file("after_kmeans_omp_centers.txt");
std::ofstream clustered_data_file("clustered_data.txt");
std::ofstream omp_clustered_data_file("omp_clustered_data.txt");

// 参数设置
const int k = 10;                  // 簇的数量
const int pointsPerCluster = 1000; // 每个簇的点数
const double stddev = 5.0;         // 生成数据的标准差
const int maxIterations = 100;     // 最大迭代次数

// 定义二维点结构
struct Point {
  double x, y;
  int cluster;

  Point(double x = 0, double y = 0, int cluster = -1)
      : x(x), y(y), cluster(cluster) {}

  // 计算两点之间的欧氏距离
  double distance(const Point &p) const {
    return std::sqrt(std::pow(x - p.x, 2) + std::pow(y - p.y, 2));
  }
};

// 生成围绕中心点的随机数据
std::vector<Point> generateData(const std::vector<Point> &centers,
                                int pointsPerCluster, double stddev) {
  std::vector<Point> data;
  std::random_device rd;
  std::mt19937 gen(rd());

  for (int i = 0; i < centers.size(); ++i) {
    std::normal_distribution<> distX(centers[i].x, stddev);
    std::normal_distribution<> distY(centers[i].y, stddev);

    // 将中心点本身添加到数据集
    data.push_back(Point(centers[i].x, centers[i].y, i));

    // 生成围绕中心点的点
    for (int j = 1; j < pointsPerCluster; ++j) {
      double x = distX(gen);
      double y = distY(gen);
      data.push_back(Point(x, y, i));
    }
  }

  return data;
}

std::vector<Point> get_random_k_center(int k) {
  std::random_device rd;
  std::mt19937 gen(rd());
  std::uniform_real_distribution<> dis(0.0, 100.0); // 正态分布

  std::vector<Point> trueCenters;
  for (int i = 0; i < k; ++i) {
    double x = dis(gen);
    double y = dis(gen);
    trueCenters.push_back(Point(x, y, i));
  }

  return trueCenters;
}

// K-means聚类算法
void kMeans(std::vector<Point> &data, std::vector<Point> &centroids, int k,
            int maxIterations) {
  bool changed = true;
  int iterations = 0;

  while (changed && iterations < maxIterations) {
    changed = false;
    iterations++;

    // 分配每个点到最近的中心点
    for (auto &point : data) {
      int oldCluster = point.cluster;
      double minDist = std::numeric_limits<double>::max();
      int minIndex = 0;

      for (int i = 0; i < k; ++i) {
        double dist = point.distance(centroids[i]);
        if (dist < minDist) {
          minDist = dist;
          minIndex = i;
        }
      }

      point.cluster = minIndex;
      if (oldCluster != point.cluster) {
        changed = true;
      }
    }

    // 如果没有点改变簇，则算法收敛
    if (!changed) {
      break;
    }

    // 重新计算中心点
    std::vector<int> counts(k, 0);
    std::vector<Point> newCentroids(k, Point(0, 0));

    for (const auto &point : data) {
      newCentroids[point.cluster].x += point.x;
      newCentroids[point.cluster].y += point.y;
      counts[point.cluster]++;
    }

    for (int i = 0; i < k; ++i) {
      if (counts[i] > 0) {
        centroids[i].x = newCentroids[i].x / counts[i];
        centroids[i].y = newCentroids[i].y / counts[i];
      }
    }
  }

  std::cout << "K-means算法在" << iterations << "次迭代后收敛" << std::endl;
}

// OpenMP版本的K-means聚类算法
void k_means_omp(std::vector<Point> &data, std::vector<Point> &centroids, int k,
                 int maxIterations, int num_threads) {
  bool changed = true;
  int iterations = 0;

  // 设置OpenMP线程数
  omp_set_num_threads(num_threads);

  while (changed && iterations < maxIterations) {
    changed = false;
    iterations++;

// 使用OpenMP并行化点分配过程
#pragma omp parallel
    {
      bool thread_changed = false;
#pragma omp for schedule(static)
      for (size_t i = 0; i < data.size(); ++i) {
        int oldCluster = data[i].cluster;
        double minDist = std::numeric_limits<double>::max();
        int minIndex = 0;

        // 计算到每个中心点的距离
        for (int j = 0; j < k; ++j) {
          double dist = data[i].distance(centroids[j]);
          if (dist < minDist) {
            minDist = dist;
            minIndex = j;
          }
        }

        data[i].cluster = minIndex;
        if (oldCluster != minIndex) {
          thread_changed = true;
        }
      }

      // 使用原子操作更新changed标志
      if (thread_changed) {
#pragma omp atomic write
        changed = true;
      }
    }

    // 如果没有点改变簇，则算法收敛
    if (!changed) {
      break;
    }

    // 并行计算新的中心点
    std::vector<std::vector<Point>> thread_centroids(
        num_threads, std::vector<Point>(k, Point(0, 0)));
    std::vector<std::vector<int>> thread_counts(num_threads,
                                                std::vector<int>(k, 0));

#pragma omp parallel
    {
      int thread_id = omp_get_thread_num();
#pragma omp for schedule(static)
      for (size_t i = 0; i < data.size(); ++i) {
        int cluster = data[i].cluster;
        thread_centroids[thread_id][cluster].x += data[i].x;
        thread_centroids[thread_id][cluster].y += data[i].y;
        thread_counts[thread_id][cluster]++;
      }
    }

    // 合并所有线程的结果
    std::vector<Point> newCentroids(k, Point(0, 0));
    std::vector<int> counts(k, 0);

    for (int t = 0; t < num_threads; ++t) {
      for (int i = 0; i < k; ++i) {
        newCentroids[i].x += thread_centroids[t][i].x;
        newCentroids[i].y += thread_centroids[t][i].y;
        counts[i] += thread_counts[t][i];
      }
    }

    // 计算新的中心点
    for (int i = 0; i < k; ++i) {
      if (counts[i] > 0) {
        centroids[i].x = newCentroids[i].x / counts[i];
        centroids[i].y = newCentroids[i].y / counts[i];
      }
    }
  }

  std::cout << "K-means算法在" << iterations << "次迭代后收敛" << std::endl;
}

void save_cluster_data(std::vector<Point> &data, std::vector<Point> &centroids,
                       std::vector<Point> &trueCenters, int op_case) {
  // 输出最终的中心点
  switch (op_case) {
  case 0:
    std::cout << "\n最终中心点(kmeans):" << std::endl;

    break;
  case 1:
    std::cout << "\n最终中心点(kmeans-omp):" << std::endl;
    break;
  case 2:
    std::cout << "\n最终中心点(kmeans-mpi):" << std::endl;
    break;
  default:
    break;
  }

  for (int i = 0; i < k; ++i) {
    std::cout << "中心点 " << i << ": (" << centroids[i].x << ", "
              << centroids[i].y << ")" << std::endl;
  }

  // 计算每个簇的点数
  std::vector<int> clusterSizes(k, 0);
  for (const auto &point : data) {
    clusterSizes[point.cluster]++;
  }

  std::cout << "\n每个簇的点数:" << std::endl;
  for (int i = 0; i < k; ++i) {
    std::cout << "簇 " << i << ": " << clusterSizes[i] << " 个点" << std::endl;
  }

  // 计算每个簇中原始类别的分布
  std::vector<std::vector<int>> clusterDistribution(k, std::vector<int>(k, 0));
  for (const auto &point : data) {
    int originalCluster = -1;
    // 找到该点的原始类别
    for (const auto &center : trueCenters) {
      if (std::abs(point.x - center.x) < 0.001 &&
          std::abs(point.y - center.y) < 0.001) {
        originalCluster = center.cluster;
        break;
      }
    }
    if (originalCluster != -1) {
      clusterDistribution[point.cluster][originalCluster]++;
    }
  }

  // 为每个新类别找到最主要的原始类别
  std::vector<int> clusterMapping(k);
  for (int i = 0; i < k; ++i) {
    int maxCount = -1;
    int dominantCluster = i;
    for (int j = 0; j < k; ++j) {
      if (clusterDistribution[i][j] > maxCount) {
        maxCount = clusterDistribution[i][j];
        dominantCluster = j;
      }
    }
    clusterMapping[i] = dominantCluster;
  }

  // 输出类别映射关系
  std::cout << "\n类别映射关系:" << std::endl;
  for (int i = 0; i < k; ++i) {
    std::cout << "新类别 " << i << " -> 原始类别 " << clusterMapping[i]
              << std::endl;
  }

  // 保存重新映射后的聚类数据点
  switch (op_case) {
  case 0:
    for (const auto &point : data) {
      clustered_data_file << point.x << "," << point.y << ","
                          << clusterMapping[point.cluster] << "\n";
    }
    break;
  case 1:
    for (const auto &point : data) {
      omp_clustered_data_file << point.x << "," << point.y << ","
                              << clusterMapping[point.cluster] << "\n";
    }
    break;
  case 2:
    break;
  default:
    break;
  }
}

int main() {

  std::random_device rd;
  std::mt19937 gen(rd());

  // 随机生成k个中心点
  std::vector<Point> trueCenters = get_random_k_center(k);

  std::cout << "生成的真实中心点:" << std::endl;
  for (int i = 0; i < k; ++i) {
    std::cout << "中心点 " << i << ": (" << trueCenters[i].x << ", "
              << trueCenters[i].y << ")" << std::endl;
  }

  // 保存真实中心点
  for (const auto &center : trueCenters) {
    true_centers_file << center.x << "," << center.y << "," << center.cluster
                      << "\n";
  }

  // 生成围绕中心点的随机数据
  std::vector<Point> data = generateData(trueCenters, pointsPerCluster, stddev);

  std::cout << "\n生成了 " << data.size() << " 个数据点" << std::endl;

  // 保存初始数据点（带有真实簇标识）
  for (const auto &point : data) {
    initial_data_file << point.x << "," << point.y << "," << point.cluster
                      << "\n";
  }

  // 打乱数据点的顺序
  std::shuffle(data.begin(), data.end(), gen);

  // 使用前k个数据点作为初始中心点
  std::vector<Point> centroids;
  for (int i = 0; i < k; ++i) {
    centroids.push_back(data[i]);
  }

  std::cout << "\n初始中心点:" << std::endl;
  for (int i = 0; i < k; ++i) {
    std::cout << "中心点 " << i << ": (" << centroids[i].x << ", "
              << centroids[i].y << ")" << std::endl;
  }

  // 运行K-means算法并计时
  auto start = std::chrono::high_resolution_clock::now();

  kMeans(data, centroids, k, maxIterations);

  auto end = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double, std::milli> duration = end - start;

  std::cout << "\nK-means串行算法运行时间: " << duration.count() << " 毫秒"
            << std::endl;
  save_cluster_data(data, centroids, trueCenters, 0);

  // 运行k-means-omp算法并计时
  int thread_num = 2;
  start = std::chrono::high_resolution_clock::now();

  k_means_omp(data, centroids, k, maxIterations, thread_num);

  end = std::chrono::high_resolution_clock::now();
  duration = end - start;

  std::cout << "\n"
            << thread_num << "线程K-means-omp算法运行时间: " << duration.count()
            << " 毫秒" << std::endl;
  save_cluster_data(data, centroids, trueCenters, 1);

  // 关闭文件
  true_centers_file.close();
  initial_data_file.close();
  clustered_data_file.close();

  return 0;
}
